import os

import bs4
from langchain_community.chat_models import ChatTongyi
from langchain_community.document_loaders import WebBaseLoader
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_chroma import Chroma
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from langchain_core.prompts import PromptTemplate
from langsmith import traceable

# 从环境变量中获取阿里云百练的 API Key
DASHSCOPE_API_KEY = os.getenv("DASHSCOPE_API_KEY")
# 阿里云百练的官网地址
DASHSCOPE_API_BASE_URL = "https://dashscope.aliyuncs.com/compatible-mode/v1"


@traceable(project_name="langchain_rag_app")
def run_rag():
    # 1. 加载文档
    # 使用 WebBaseLoader 从网页加载内容，并仅保留标题、标题头和文章内容
    bs4_strainer = bs4.SoupStrainer(class_=("post-title", "post-header", "post-content"))
    loader = WebBaseLoader(
        web_paths=("https://lilianweng.github.io/posts/2025-05-01-thinking/",),
        bs_kwargs={"parse_only": bs4_strainer},
    )
    docs = loader.load()

    # 2. 文档分割
    # 使用 RecursiveCharacterTextSplitter 将文档分割成块，每块1000字符，重叠200字符
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000, chunk_overlap=200, add_start_index=True
    )
    all_splits = text_splitter.split_documents(docs)

    # 3. 存储嵌入
    # 使用 Chroma 向量存储和 DashScopeEmbeddings(model="text-embedding-v1") 模型，将分割的文档块嵌入并存储
    vector_store = Chroma.from_documents(
        documents=all_splits,
        embedding=DashScopeEmbeddings(model="text-embedding-v1")
    )

    # 4. 检索文档
    # 使用 VectorStoreRetriever 从向量存储中检索与查询最相关的文档
    retriever = vector_store.as_retriever(search_type="similarity", search_kwargs={"k": 3})

    # 5. 生成回答
    llm = ChatTongyi(model="qwen-max")
    # 自定义提示词模板
    template = """Use the following pieces of context to answer the question at the end.
    If you don't know the answer, just say that you don't know, don't try to make up an answer.
    Use three sentences maximum and keep the answer as concise as possible.
    Always say "thanks for asking!" at the end of the answer.
    
    {context}
    
    Question: {question}
    
    Helpful Answer:"""
    rag_prompt = PromptTemplate.from_template(template)

    # 定义格式化文档的函数
    def format_docs(docs):
        return "\n\n".join(doc.page_content for doc in docs)

    # 使用 LangChain Expression Language (LCEL) 构建 RAG Chain
    rag_chain = (
            {"context": retriever | format_docs, "question": RunnablePassthrough()}
            | rag_prompt
            | llm
            | StrOutputParser()
    )

    # 流式生成回答
    for chunk in rag_chain.stream("Why We Think?"):
        print(chunk, end="", flush=True)


if __name__ == '__main__':
    run_rag()


